package org.example.myleet.p407;

import java.util.*;

public class Solution {

    /**
     * row, col change
     */
    private static final int[][] DIRECTIONS = new int[][]{{-1, 0}, {0, 1}, {1, 0}, {0, -1}};

    public int trapRainWater(int[][] heightMap) {
        int m = heightMap.length, n = heightMap[0].length;
        int[][] water = new int[m][n];
        int maxHeight = 0;
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                maxHeight = Math.max(maxHeight, heightMap[i][j]);
            }
        }
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                water[i][j] = maxHeight;
            }
        }
        for (int j = 0; j < n; ++j) {
            water[0][j] = heightMap[0][j];
            water[m - 1][j] = heightMap[m - 1][j];
        }
        for (int i = 0; i < m; ++i) {
            water[i][0] = heightMap[i][0];
            water[i][n - 1] = heightMap[i][n - 1];
        }
        Queue<int[]> queue = new LinkedList<>();
        init:
        for (int i = 1; i < m - 1; ++i) {
            for (int j = 1; j < n - 1; ++j) {
                List<Integer> dirs = new ArrayList<>(4);
                for (int k = 0; k < 4; ++k) {
                    int[] dir = DIRECTIONS[k];
                    dirs.add(water[i + dir[0]][j + dir[1]]);
                }
                dirs.sort(Comparator.naturalOrder());
                int newWaterHeight = Math.max(heightMap[i][j], dirs.get(0));
                if (water[i][j] > newWaterHeight) {
                    water[i][j] = newWaterHeight;
                    for (int k = 0; k < 4; ++k) {
                        int[] dir = DIRECTIONS[k];
                        int nextRow = i + dir[0];
                        int nextCol = j + dir[1];
                        if (0 < nextRow && nextRow < m - 1 && 0 < nextCol && nextCol < n - 1) {
                            queue.offer(new int[]{nextRow, nextCol});
                        }
                    }
                    break init;
                }
            }
        }
        while (!queue.isEmpty()) {
            int[] next = queue.poll();
            List<Integer> dirs = new ArrayList<>(4);
            for (int k = 0; k < 4; ++k) {
                int[] dir = DIRECTIONS[k];
                dirs.add(water[next[0] + dir[0]][next[1] + dir[1]]);
            }
            dirs.sort(Comparator.naturalOrder());
            int newWaterHeight = Math.max(heightMap[next[0]][next[1]], dirs.get(0));
            if (water[next[0]][next[1]] > newWaterHeight) {
                water[next[0]][next[1]] = newWaterHeight;
                for (int k = 0; k < 4; ++k) {
                    int[] dir = DIRECTIONS[k];
                    int nextRow = next[0] + dir[0];
                    int nextCol = next[1] + dir[1];
                    if (0 < nextRow && nextRow < m - 1 && 0 < nextCol && nextCol < n - 1) {
                        queue.offer(new int[]{nextRow, nextCol});
                    }
                }
            }
        }
        int rain = 0;
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                rain += water[i][j] - heightMap[i][j];
            }
        }
        return rain;
    }
}
